# Copyright © 2023 Pathway

from __future__ import annotations

import datetime
import warnings
from collections.abc import Iterable
from dataclasses import dataclass
from types import EllipsisType
from typing import TYPE_CHECKING, Any, TypeVar

from pathway.internals import dtype as dt, expression as expr
from pathway.internals.expression_printer import get_expression_info
from pathway.internals.expression_visitor import IdentityTransform
from pathway.internals.json import Json
from pathway.internals.operator_mapping import (
    common_dtype_in_binary_operator,
    get_binary_operators_mapping,
    get_binary_operators_mapping_optionals,
    get_unary_operators_mapping,
    tuple_handling_operators,
)

if TYPE_CHECKING:
    from pathway.internals.table import Table


@dataclass(eq=True, frozen=True)
class TypeInterpreterState:
    unoptionalize: Iterable[expr.InternalColRef] = ()

    def with_new_col(
        self, new_cols: Iterable[expr.ColumnReference]
    ) -> TypeInterpreterState:
        return TypeInterpreterState(
            [
                *self.unoptionalize,
                *[col._to_internal() for col in new_cols],
            ]
        )

    def check_colref_to_unoptionalize_from_colrefs(
        self, colref: expr.ColumnReference
    ) -> bool:
        return colref._to_internal() in self.unoptionalize

    def check_colref_to_unoptionalize_from_tables(
        self, colref: expr.ColumnReference
    ) -> bool:
        return self.check_colref_to_unoptionalize_from_colrefs(colref.table.id)


class TypeInterpreter(IdentityTransform):
    def eval_column_val(
        self,
        expression: expr.ColumnReference,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.ColumnReference:
        expression = super().eval_column_val(expression, state=state, **kwargs)
        return _wrap(expression, self._eval_column_val(expression, state))

    def _eval_column_val(
        self,
        expression: expr.ColumnReference,
        state: TypeInterpreterState | None = None,
    ) -> dt.DType:
        from pathway.internals.table import Table

        assert isinstance(expression._table, Table)
        dtype = expression._column.dtype
        assert state is not None
        assert isinstance(dtype, dt.DType)
        if state.check_colref_to_unoptionalize_from_colrefs(expression):
            return dt.unoptionalize(dtype)
        return dtype

    def eval_unary_op(
        self,
        expression: expr.ColumnUnaryOpExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.ColumnUnaryOpExpression:
        expression = super().eval_unary_op(expression, state=state, **kwargs)
        operand_dtype = expression._expr._dtype
        operator_fun = expression._operator
        if (
            dtype := get_unary_operators_mapping(operator_fun, operand_dtype)
        ) is not None:
            return _wrap(expression, dtype)
        expression_info = get_expression_info(expression)
        raise TypeError(
            f"Pathway does not support using unary operator {operator_fun.__name__}"
            + f" on column of type {expression._dtype.typehint}.\n"
            + "It refers to the following expression:\n"
            + expression_info
        )

    def eval_binary_op(
        self,
        expression: expr.ColumnBinaryOpExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.ColumnBinaryOpExpression:
        expression = super().eval_binary_op(expression, state=state, **kwargs)
        left_dtype = expression._left._dtype
        right_dtype = expression._right._dtype
        return _wrap(
            expression,
            self._eval_binary_op(
                left_dtype, right_dtype, expression._operator, expression
            ),
        )

    def _eval_binary_op(
        self,
        left_dtype: dt.DType,
        right_dtype: dt.DType,
        operator: Any,
        expression: expr.ColumnExpression,
    ) -> dt.DType:
        original_left = left_dtype
        original_right = right_dtype
        if (
            dtype := common_dtype_in_binary_operator(left_dtype, right_dtype)
        ) is not None:
            left_dtype = dtype
            right_dtype = dtype

        if (
            dtype := get_binary_operators_mapping(operator, left_dtype, right_dtype)
        ) is not None:
            return dtype

        left_dtype, right_dtype = dt.unoptionalize_pair(left_dtype, right_dtype)

        if (
            dtype_and_handler := get_binary_operators_mapping_optionals(
                operator, left_dtype, right_dtype
            )
        ) is not None:
            return dtype_and_handler[0]

        if (
            isinstance(left_dtype, dt.Tuple)
            and isinstance(right_dtype, dt.Tuple)
            and operator in tuple_handling_operators
        ):
            left_args = left_dtype.args
            right_args = right_dtype.args
            if (
                isinstance(left_args, tuple)
                and isinstance(right_args, tuple)
                and len(left_args) == len(right_args)
            ):
                results = tuple(
                    self._eval_binary_op(left_arg, right_arg, operator, expression)
                    for left_arg, right_arg in zip(left_args, right_args)
                    if not isinstance(left_arg, EllipsisType)
                    and not isinstance(right_arg, EllipsisType)
                )
                assert all(result == results[0] for result in results)
                return dt.wrap(results[0])
        expression_info = get_expression_info(expression)
        raise TypeError(
            f"Pathway does not support using binary operator {operator.__name__}"
            + f" on columns of types {original_left.typehint}, {original_right.typehint}.\n"
            + "It refers to the following expression:\n"
            + expression_info
        )

    def eval_const(
        self,
        expression: expr.ColumnConstExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.ColumnConstExpression:
        expression = super().eval_const(expression, state=state, **kwargs)
        type_ = type(expression._val)
        if isinstance(expression._val, datetime.datetime):
            if expression._val.tzinfo is None:
                dtype: dt.DType = dt.DATE_TIME_NAIVE
            else:
                dtype = dt.DATE_TIME_UTC
        elif isinstance(expression._val, datetime.timedelta):
            dtype = dt.DURATION
        else:
            dtype = dt.wrap(type_)
        return _wrap(expression, dtype)

    def eval_reducer(
        self,
        expression: expr.ReducerExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.ReducerExpression:
        expression = super().eval_reducer(expression, state=state, **kwargs)
        return _wrap(expression, self._eval_reducer(expression))

    def _eval_reducer(
        self,
        expression: expr.ReducerExpression,
    ) -> dt.DType:
        return expression._reducer.return_type([arg._dtype for arg in expression._args])

    def eval_count(
        self,
        expression: expr.CountExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.CountExpression:
        expression = super().eval_count(expression, state=state, **kwargs)
        return _wrap(expression, dt.INT)

    def eval_apply(
        self,
        expression: expr.ApplyExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.ApplyExpression:
        expression = super().eval_apply(expression, state=state, **kwargs)
        return _wrap(expression, expression._return_type)

    def eval_numbaapply(
        self,
        expression: expr.NumbaApplyExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.NumbaApplyExpression:
        expression = super().eval_numbaapply(expression, state=state, **kwargs)
        return _wrap(expression, expression._return_type)

    def eval_async_apply(
        self,
        expression: expr.AsyncApplyExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.AsyncApplyExpression:
        expression = super().eval_async_apply(expression, state=state, **kwargs)
        return _wrap(expression, expression._return_type)

    def eval_call(
        self,
        expression: expr.ColumnCallExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.ColumnCallExpression:
        dtype = expression._col_expr._column.dtype
        assert isinstance(dtype, dt.Callable)
        expression = super().eval_call(expression, state=state, **kwargs)
        arg_annots, ret_type = dtype.arg_types, dtype.return_type
        if not isinstance(arg_annots, EllipsisType):
            arg_annots = arg_annots[1:]  # ignoring self
            for arg_annot, arg in zip(arg_annots, expression._args, strict=True):
                arg_dtype = arg._dtype
                assert not isinstance(arg_annot, EllipsisType)
                assert dt.dtype_issubclass(arg_dtype, arg_annot)
        return _wrap(expression, ret_type)

    def eval_pointer(
        self,
        expression: expr.PointerExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.PointerExpression:
        expression = super().eval_pointer(expression, state=state, **kwargs)
        arg_types = [arg._dtype for arg in expression._args]
        if expression._optional and any(
            isinstance(arg, dt.Optional) for arg in arg_types
        ):
            return _wrap(expression, dt.Optional(dt.POINTER))
        else:
            return _wrap(expression, dt.POINTER)

    def eval_cast(
        self,
        expression: expr.CastExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.CastExpression:
        expression = super().eval_cast(expression, state=state, **kwargs)
        return _wrap(expression, expression._return_type)

    def eval_convert(
        self,
        expression: expr.ConvertExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.ConvertExpression:
        expression = super().eval_convert(expression, state=state, **kwargs)
        return _wrap(expression, expression._return_type)

    def eval_declare(
        self,
        expression: expr.DeclareTypeExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.DeclareTypeExpression:
        expression = super().eval_declare(expression, state=state, **kwargs)
        return _wrap(expression, expression._return_type)

    def eval_coalesce(
        self,
        expression: expr.CoalesceExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.CoalesceExpression:
        expression = super().eval_coalesce(expression, state=state, **kwargs)
        dtypes = [arg._dtype for arg in expression._args]
        ret_type = dtypes[0]
        non_optional_arg = False
        for dtype in dtypes:
            ret_type = dt.types_lca(dtype, ret_type)
            if not isinstance(dtype, dt.Optional):
                # FIXME: do we want to be more radical and return now?
                # Maybe with a warning that some args are skipped?
                non_optional_arg = True
        if ret_type is dt.ANY and any(dtype is not dt.ANY for dtype in dtypes):
            raise TypeError(
                f"Cannot perform pathway.coalesce on columns of types {[dtype.typehint for dtype in dtypes]}."
            )
        ret_type = dt.unoptionalize(ret_type) if non_optional_arg else ret_type
        return _wrap(expression, ret_type)

    def eval_require(
        self,
        expression: expr.RequireExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.RequireExpression:
        assert state is not None
        args = [
            self.eval_expression(arg, state=state, **kwargs) for arg in expression._args
        ]
        new_state = state.with_new_col(
            [arg for arg in expression._args if isinstance(arg, expr.ColumnReference)]
        )
        val = self.eval_expression(expression._val, state=new_state, **kwargs)
        expression = expr.RequireExpression(val, *args)
        ret_type = dt.Optional(val._dtype)
        return _wrap(expression, ret_type)

    def eval_not_none(
        self,
        expression: expr.IsNotNoneExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.IsNotNoneExpression:
        ret = super().eval_not_none(expression, state=state, **kwargs)
        return _wrap(ret, dt.BOOL)

    def eval_none(
        self,
        expression: expr.IsNoneExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.IsNoneExpression:
        ret = super().eval_none(expression, state=state, **kwargs)
        return _wrap(ret, dt.BOOL)

    def eval_ifelse(
        self,
        expression: expr.IfElseExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.IfElseExpression:
        expression = super().eval_ifelse(expression, state=state, **kwargs)
        if_dtype = expression._if._dtype
        if if_dtype != dt.BOOL:
            raise TypeError(
                f"First argument of pathway.if_else has to be bool, found {if_dtype.typehint}."
            )
        then_dtype = expression._then._dtype
        else_dtype = expression._else._dtype
        lca = dt.types_lca(then_dtype, else_dtype)
        if lca is dt.ANY:
            raise TypeError(
                f"Cannot perform pathway.if_else on columns of types {then_dtype.typehint} and {else_dtype.typehint}."
            )
        return _wrap(expression, lca)

    def eval_make_tuple(
        self,
        expression: expr.MakeTupleExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.MakeTupleExpression:
        expression = super().eval_make_tuple(expression, state=state, **kwargs)
        dtypes = tuple(arg._dtype for arg in expression._args)
        return _wrap(expression, dt.Tuple(*dtypes))

    def _eval_json_get(
        self,
        expression: expr.GetExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.GetExpression:
        return _wrap(expression, dt.JSON)

    def eval_get(
        self,
        expression: expr.GetExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.GetExpression:
        expression = super().eval_get(expression, state=state, **kwargs)
        object_dtype = expression._object._dtype
        index_dtype = expression._index._dtype
        default_dtype = expression._default._dtype

        if object_dtype == dt.JSON:
            # json
            if not default_dtype.is_subclass_of(dt.Optional(dt.JSON)):
                raise TypeError(
                    f"Default must be of type {Json | None}, found {default_dtype.typehint}."
                )
            if not expression._check_if_exists or default_dtype == dt.JSON:
                return _wrap(expression, dt.JSON)
            else:
                return _wrap(expression, dt.Optional(dt.JSON))
        elif object_dtype.equivalent_to(dt.Optional(dt.JSON)):
            # optional json
            raise TypeError(f"Cannot get from {Json | None}.")
        else:
            # sequence
            if not isinstance(
                object_dtype, (dt.Tuple, dt.List)
            ) and object_dtype not in [
                dt.ANY,
                dt.ARRAY,
            ]:
                raise TypeError(
                    f"Object in {expression!r} has to be a JSON or sequence."
                )
            if index_dtype != dt.INT:
                raise TypeError(f"Index in {expression!r} has to be an int.")

            if object_dtype == dt.ARRAY:
                warnings.warn(
                    f"Object in {expression!r} is of type numpy.ndarray but its number of"
                    + " dimensions is not known. Pathway cannot determine the return type"
                    + " and will set Any as the return type. Please use "
                    + "pathway.declare_type to set the correct return type."
                )
                return _wrap(expression, dt.ANY)
            if object_dtype == dt.ANY:
                return _wrap(expression, dt.ANY)

            if isinstance(object_dtype, dt.List):
                if expression._check_if_exists:
                    return _wrap(expression, dt.Optional(object_dtype.wrapped))
                else:
                    return _wrap(expression, object_dtype.wrapped)
            assert isinstance(object_dtype, dt.Tuple)
            if object_dtype == dt.ANY_TUPLE:
                return _wrap(expression, dt.ANY)

            assert not isinstance(object_dtype.args, EllipsisType)
            dtypes = object_dtype.args

            if (
                expression._const_index is None
            ):  # no specified position, index is an Expression
                assert isinstance(dtypes[0], dt.DType)
                return_dtype = dtypes[0]
                for dtype in dtypes[1:]:
                    if isinstance(dtype, dt.DType):
                        return_dtype = dt.types_lca(return_dtype, dtype)
                if expression._check_if_exists:
                    return_dtype = dt.types_lca(return_dtype, default_dtype)
                return _wrap(expression, return_dtype)

            if not isinstance(expression._const_index, int):
                raise IndexError("Index n")

            try:
                try_ret = dtypes[expression._const_index]
                return _wrap(expression, try_ret)
            except IndexError:
                message = (
                    f"Index {expression._const_index} out of range for a tuple of"
                    + f" type {object_dtype.typehint}."
                )
                if expression._check_if_exists:
                    expression_info = get_expression_info(expression)
                    warnings.warn(
                        message
                        + " It refers to the following expression:\n"
                        + expression_info
                        + "Consider using just the default value without .get()."
                    )
                    return _wrap(expression, default_dtype)
                else:
                    raise IndexError(message)

    def eval_method_call(
        self,
        expression: expr.MethodCallExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.MethodCallExpression:
        expression = super().eval_method_call(expression, state=state, **kwargs)
        dtypes = tuple([arg._dtype for arg in expression._args])
        if (dtypes_and_handler := expression.get_function(dtypes)) is not None:
            return _wrap(expression, dt.wrap(dtypes_and_handler[1]))

        if len(dtypes) > 0:
            with_arguments = (
                f" with arguments of type {[dtype.typehint for dtype in dtypes[1:]]}"
            )
        else:
            with_arguments = ""
        raise AttributeError(
            f"Column of type {dtypes[0].typehint} has no attribute {expression._name}{with_arguments}."
        )

    def eval_unwrap(
        self,
        expression: expr.UnwrapExpression,
        state: TypeInterpreterState | None = None,
        **kwargs,
    ) -> expr.UnwrapExpression:
        expression = super().eval_unwrap(expression, state=state, **kwargs)
        dtype = expression._expr._dtype
        return _wrap(expression, dt.unoptionalize(dtype))


class JoinTypeInterpreter(TypeInterpreter):
    left: Table
    right: Table
    optionalize_left: bool
    optionalize_right: bool

    def __init__(self, left, right, optionalize_left, optionalize_right):
        self.left = left
        self.right = right
        self.optionalize_left = optionalize_left
        self.optionalize_right = optionalize_right
        super().__init__()

    def _eval_column_val(
        self,
        expression: expr.ColumnReference,
        state: TypeInterpreterState | None = None,
    ) -> dt.DType:
        dtype = expression._column.dtype
        assert state is not None
        if (expression.table == self.left and self.optionalize_left) or (
            expression.table == self.right and self.optionalize_right
        ):
            return dt.Optional(dtype)
        return super()._eval_column_val(expression, state=state)


class JoinRowwiseTypeInterpreter(TypeInterpreter):
    temporary_column_to_original: dict[expr.InternalColRef, expr.InternalColRef]
    original_column_to_temporary: dict[expr.InternalColRef, expr.ColumnReference]

    def __init__(
        self,
        temporary_column_to_original: dict[expr.InternalColRef, expr.InternalColRef],
        original_column_to_temporary: dict[expr.InternalColRef, expr.ColumnReference],
    ) -> None:
        self.temporary_column_to_original = temporary_column_to_original
        self.original_column_to_temporary = original_column_to_temporary
        super().__init__()

    def _eval_column_val(
        self,
        expression: expr.ColumnReference,
        state: TypeInterpreterState | None = None,
    ) -> dt.DType:
        assert state is not None
        if expression._to_internal() in self.temporary_column_to_original:
            # if clause needed because of external columns from ix() that are not in
            # self.temporary_column_to_original
            original = self.temporary_column_to_original[expression._to_internal()]
            tmp_id_colref = self.original_column_to_temporary[
                original._table.id._to_internal()
            ]
            if state.check_colref_to_unoptionalize_from_colrefs(tmp_id_colref):
                return original._column.dtype
        return super()._eval_column_val(expression, state=state)


def eval_type(expression: expr.ColumnExpression) -> dt.DType:
    return (
        TypeInterpreter()
        .eval_expression(expression, state=TypeInterpreterState())
        ._dtype
    )


ColExprT = TypeVar("ColExprT", bound=expr.ColumnExpression)


def _wrap(expression: ColExprT, dtype: dt.DType) -> ColExprT:
    assert not hasattr(expression, "_dtype")
    assert isinstance(dtype, dt.DType)
    expression._dtype = dtype
    return expression
